// clang-format off
//#include <Python.h>
// clang-format on
#include "DoseStatTorch.h"

#include <ATen/Operators.h>
#include <torch/all.h>
#include <torch/library.h>
#include <torch/types.h>

#include <algorithm>
#include <stdexcept>

#include "c10/core/Scalar.h"

///xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx 控制是否打印
#if 0
#    define MYDBG printf
#else
#    define MYDBG(...)
#endif

extern "C"
{
    /* Creates a dummy empty _C module that can be imported from Python.
The import from Python will load the .so consisting of this file
in this extension, so that the TORCH_LIBRARY static initializers
below are run. */
    PyObject* PyInit_dose_torch(void)
    {
        static struct PyModuleDef module_def = {
            PyModuleDef_HEAD_INIT,
            "dose_torch", /* name of module */
            NULL,         /* module documentation, may be NULL */
            -1,           /* size of per-interpreter state of the module,
               or -1 if the module keeps state in global variables. */
            NULL,         /* methods */
        };
        return PyModule_Create(&module_def);
    }
}

static torch::Tensor FromVectorCopy(const std::vector<double>& v)
{
    const int64_t n = static_cast<int64_t>(v.size());
    // 先造一个空的 Tensor，再把数据 memcpy 进去
    torch::Tensor t = torch::zeros({n}, torch::kFloat64);
    std::memcpy(t.data_ptr<double>(), v.data(), n * sizeof(double));
    return t;
}

namespace dose_torch
{
    // 计算DVH(剂量体积直方图)的PyTorch自定义算子实现
    std::tuple<torch::Tensor, torch::Tensor> accumulate(const torch::Tensor& voxelVolume, const torch::Tensor& voxelDose, bool doExactAccumulation)
    {
        // 检查输入张量的形状和类型
        TORCH_CHECK(voxelVolume.dim() == 1, "voxelVolume must be 1D tensor");
        TORCH_CHECK(voxelDose.dim() == 1, "voxelDose must be 1D tensor");
        TORCH_CHECK(voxelVolume.size(0) == voxelDose.size(0), "voxelVolume and voxelDose must have the same length");
        TORCH_CHECK(voxelVolume.scalar_type() == torch::kFloat64, "voxelVolume must be double tensor");
        TORCH_CHECK(voxelDose.scalar_type() == torch::kFloat64, "voxelDose must be double tensor");

        const int64_t numVoxels = voxelVolume.size(0);
        auto voxelVolume_acc = voxelVolume.accessor<double, 1>();
        auto voxelDose_acc = voxelDose.accessor<double, 1>();

        MYDBG("doExactAccumulation: %d\n", doExactAccumulation);

        // 预分配可能的内存以减少动态扩容; dose->volume
        std::vector<std::pair<double, double> > doseVolumePairs;
        doseVolumePairs.reserve(numVoxels);

        /// 过滤体积为0的体素, 同时直接创建 剂量->体积 对
        for (int64_t i = 0; i < numVoxels; ++i)
        {
            if (voxelVolume_acc[i] > 0)
            {
                doseVolumePairs.emplace_back(std::max(0.0, voxelDose_acc[i]), voxelVolume_acc[i]);
            }
        }
        doseVolumePairs.shrink_to_fit();
        // number of positive voxels
        const size_t num_pstv = doseVolumePairs.size();

        ///==== 空数据的快速处理
        if (num_pstv == 0)
        {
            /// ticks 最少为 2 => 一个 bin
            /// volume; 首哨兵为总体积, 末哨兵为 0.0
            /// dose; 首哨兵为 min dose, 末哨兵为 max dose
#if 0
            torch::Tensor result = torch::zeros({2, 2}, torch::kFloat64);
            auto result_acc = result.accessor<double, 2>();
            result_acc[0][0] = 1.0;  // accumulatedRelativeVolumes[0]
            result_acc[0][1] = 0.0;  // accumulatedRelativeVolumes[1]
            result_acc[1][0] = 0.0;  // binnedDoseValues[0]
            result_acc[1][1] = 0.0;  // binnedDoseValues[1]
            return result;
#endif
            torch::Tensor relativeVolTsr = torch::tensor({1.0, 0.0}, torch::kFloat64);
            torch::Tensor binDoseTsr = torch::tensor({0.0, 0.0}, torch::kFloat64);

            MYDBG("Warning: Empty input vectors provided to accumulate\n");
            return {binDoseTsr, relativeVolTsr};
        }

        /// 根据数据量和doExactAccumulation参数选择合适的算法
        constexpr size_t EXACT_THRESHOLD = 5000;
        constexpr int MAX_BINS = 1001;
        constexpr int MIN_BINS = 2;
        constexpr double doseMinWidth = 1e-6;  // dose bin 最小分辨率
        const bool useExactMethod = doExactAccumulation || (num_pstv < EXACT_THRESHOLD);

        std::vector<double> binnedDoseValues;
        std::vector<double> accumulatedRelativeVolumes;

        do
        {
            ///============================= 精确计算:
            /// 直接排序 {剂量->体积}, 先按照剂量排序, 然后按照体积排序
            if (useExactMethod)
            {
                // 避免额外的索引数组和多次拷贝
                std::sort(doseVolumePairs.begin(), doseVolumePairs.end());

                // 准备结果数组
                binnedDoseValues.resize(num_pstv + 2, 0.0);
                accumulatedRelativeVolumes.resize(num_pstv + 2, 0.0);

                // 填充排序后的结果
                for (size_t i = 0; i < num_pstv; ++i)
                {
                    binnedDoseValues[i + 1] = doseVolumePairs[i].first;
                    accumulatedRelativeVolumes[i + 1] = doseVolumePairs[i].second;
                }

                break;  // skip all the rest
            }

            ///============================= 优化的粗略计算:
            /// 基于统计特征选择更合理的分箱数量; 提取剂量数据进行统计分析
            std::vector<double> doses;
            doses.reserve(num_pstv);
            double totalVolume = 0.0;

            for (const auto& [dose, volume] : doseVolumePairs)
            {
                doses.push_back(dose);
                totalVolume += volume;
            }

            /// 计算剂量范围
            auto [minDoseIt, maxDoseIt] = std::minmax_element(doses.begin(), doses.end());
            double minDose = std::max(0.0, *minDoseIt);
            double maxDose = *maxDoseIt;
            MYDBG("[my]minDose: %f, maxDose: %f, num_pstv: %zu\n", minDose, maxDose, num_pstv);

            /// 动态计算合适的分箱数
            const double doseRange = maxDose - minDose;

            ///============================= 粗略; 当剂量范围<=最小宽度, 用最小分箱数
            if (doseRange <= doseMinWidth)
            {
                char errMsg[256];
                sprintf(errMsg,
                        "Error: dose range to small, range -> %f, min -> %f, dose -> %f; while voxels number too much -> (%zu); check input data.\n",
                        doseRange, minDose, maxDose, num_pstv);
                throw std::invalid_argument(errMsg);
            }

            ///============================= 当剂量范围>最小宽度, 动态计算分箱数
            // 根据剂量范围和体素数量自适应确定分箱数
            int numBins = MIN_BINS;
            // 计算理论分箱数: 每doseMinWidth剂量一个分箱; 加1.0 确保覆盖整个剂量范围
            const double theoreticalBins = doseRange / doseMinWidth + 1.0;
            // 将理论分箱数转换为整数, 并限制在最小和最大分箱数之间
            const auto calculatedBins = static_cast<int>(theoreticalBins);
            numBins = std::max(MIN_BINS, std::min(MAX_BINS, calculatedBins));

            /// 准备分箱结果
            // 0 位置, n+1 位置, 两个哨兵
            constexpr size_t tickOfst = 1;
            binnedDoseValues.resize(numBins + 2, 0.0);
            accumulatedRelativeVolumes.resize(numBins + 2, 0.0);
            const double binWidth = doseRange / (numBins - 1);

            // 填充分箱剂量值
            for (int j = 0; j < numBins; ++j)
            {
                binnedDoseValues[j + tickOfst] = minDose + j * binWidth;
            }

            /// 统计每个分箱的体积
            for (const auto& [dose, volume] : doseVolumePairs)
            {
                if (dose <= minDose)
                {
                    // 哨兵 0
                    accumulatedRelativeVolumes[tickOfst] += volume;
                    continue;
                }
                if (dose >= maxDose)
                {
                    // 哨兵 n+1
                    continue;
                }
                // 落点取 bin 的右边界, 对应高的 dose
                const int binIndex = static_cast<int>((dose - minDose) / binWidth + 0.5);
                accumulatedRelativeVolumes[binIndex + tickOfst] += volume;
            }

            break;
        } while (0);

        // 计算累积相对体积
        const auto numBins = static_cast<int>(accumulatedRelativeVolumes.size()) - 2;
        double totalVol = 0.0;
        for (int i = numBins; i > 0; --i)
        {
            /// i-1 的累积体积 + 0.5 的当前体积, 偏低; counter 上面取右侧偏高;
            const double vol = accumulatedRelativeVolumes[i];
            accumulatedRelativeVolumes[i] = totalVol + 0.5 * vol;
            totalVol += vol;
        }

        /// 哨兵 规整
        accumulatedRelativeVolumes[0] = totalVol;
        accumulatedRelativeVolumes.back() = 0.0;
        binnedDoseValues[0] = binnedDoseValues[1];
        binnedDoseValues.back() = binnedDoseValues[numBins];

        /// 归一化
        if (totalVol > 0.0)
        {
            for (auto& volume : accumulatedRelativeVolumes)
            {
                volume /= totalVol;
            }
        }

        // 将结果转换为PyTorch张量
#if 0
        const int64_t result_size = binnedDoseValues.size();
        torch::Tensor result = torch::zeros({2, result_size}, torch::kFloat64);
        auto result_acc = result.accessor<double, 2>();

        for (int64_t i = 0; i < result_size; ++i)
        {
            result_acc[0][i] = accumulatedRelativeVolumes[i];
            result_acc[1][i] = binnedDoseValues[i];
        }
        return result;
#else
        auto binDoseValuesTsr = torch::tensor(binnedDoseValues, torch::kF64);
        auto accumReltvVolumesTsr = torch::tensor(accumulatedRelativeVolumes, torch::kF64);

        return {binDoseValuesTsr, accumReltvVolumesTsr};
#endif
    }

    // 获取特定体积百分比对应的剂量的PyTorch自定义算子实现
    torch::Scalar get_metric(const torch::Tensor& binnedDoseValues, const torch::Tensor& accumulatedRelativeVolumes, double inDoseAtVolume)
    {
        // 检查输入张量的形状和类型
        //TORCH_CHECK(inDoseAtVolume.dim() == 0, "inDoseAtVolume must be scalar tensor");
        TORCH_CHECK(binnedDoseValues.dim() == 1, "binnedDoseValues must be 1D tensor");
        TORCH_CHECK(accumulatedRelativeVolumes.dim() == 1, "accumulatedRelativeVolumes must be 1D tensor");
        TORCH_CHECK(binnedDoseValues.size(0) == accumulatedRelativeVolumes.size(0),
                    "binnedDoseValues and accumulatedRelativeVolumes must have the same length");

        /// declare scalar type, dimension
        //const auto doseAtVolume_acc = *inDoseAtVolume.data_ptr<double>();
        //const double doseAtVolume_val = doseAtVolume_acc;
        const double doseAtVolume_val = inDoseAtVolume;
        auto binDoseValAcsr = binnedDoseValues.accessor<double, 1>();
        auto accRelVolumesAcsr = accumulatedRelativeVolumes.accessor<double, 1>();

        const int64_t doseSz = binnedDoseValues.size(0);

        /// 参数校验
        if (doseSz == 0)
        {
            MYDBG("Warning: Empty input vectors provided to GetMetric\n");
            return torch::Scalar(0.0);
        }

        /// 边界检查优化
        if (doseAtVolume_val >= accRelVolumesAcsr[0])
        {
            return torch::Scalar(binDoseValAcsr[0]);
        }
        if (doseAtVolume_val <= accRelVolumesAcsr[doseSz - 1])
        {
            return torch::Scalar(binDoseValAcsr[doseSz - 1]);
        }

        /// 二分查找, 返回 所有等于 value 的元素范围
        std::vector<double> accumulatedRelativeVolumes_vec(accRelVolumesAcsr.data(), accRelVolumesAcsr.data() + doseSz);
        auto [first, last] =
            std::equal_range(accumulatedRelativeVolumes_vec.begin(), accumulatedRelativeVolumes_vec.end(), doseAtVolume_val, std::greater<double>{});

        ///================= 存在恰好等于 dose 的情况 [v1,v2,...], rarely
        if (first != last)
        {
            // 计算first到last范围内binnedDoseValues的平均值
            size_t count = 0;
            double sumDoses = 0.0;
            for (auto it = first; it != last; ++it)
            {
                const size_t index = it - accumulatedRelativeVolumes_vec.begin();
                sumDoses += binDoseValAcsr[index];
                count++;
            }
            const double avgDose = sumDoses / count;
            MYDBG("Volume -> %f, matching values -> (count: %zu), average dose -> %f \n", doseAtVolume_val, count, avgDose);
            return torch::Scalar(avgDose);
        }

        ///================= 不存在等于 dose 的情况, 插值; lower_bound == upper_bound
        const size_t ind = first - accumulatedRelativeVolumes_vec.begin();
        MYDBG("Target dose -> %f, found index -> %zu\n", doseAtVolume_val, ind);
        ///------------- 边界检查
        if (ind == 0)
        {
            return torch::Scalar(binDoseValAcsr[0]);
        }
        if (ind >= doseSz)
        {
            return torch::Scalar(binDoseValAcsr[doseSz - 1]);
        }
        size_t ind1 = ind - 1;  // little dose; big volume
        size_t ind2 = ind;      // big dose; small volume

        ///------------- 如果体积接近, 返回平均剂量
        const double volDiff = accRelVolumesAcsr[ind2] - accRelVolumesAcsr[ind1];
        if (std::abs(volDiff) <= 1e-10)
        {
            MYDBG("Average; dose -> %f, index -> %zu\n", doseAtVolume_val, ind);
            return torch::Scalar(0.5 * (binDoseValAcsr[ind2] + binDoseValAcsr[ind1]));
        }

        ///------------- 线性插值
        MYDBG("Linear interpolation; dose -> %f, index -> %zu\n", doseAtVolume_val, ind);
        double doseDiff = binDoseValAcsr[ind2] - binDoseValAcsr[ind1];
        double metric = binDoseValAcsr[ind1] + (doseAtVolume_val - accRelVolumesAcsr[ind1]) * doseDiff / volDiff;
        return torch::Scalar(metric);
    }

    // 计算ROI内的剂量统计指标的PyTorch自定义算子实现
    torch::Scalar get_dose_stat_roi_interp(const torch::Tensor& binnedDoseValues,
                                           const torch::Tensor& accumulatedRelativeVolumes,
                                           const torch::Tensor& voxelVolume,
                                           const torch::Tensor& voxelDose,
                                           std::string_view dose_type,
                                           double doseAtVolume)
    {
        ///================= 计算平均剂量时, 可以提前检查数组大小是否匹配
        if (voxelVolume.size(0) != voxelDose.size(0))
        {
            throw std::runtime_error("voxelVolume and voxelDose arrays have different sizes");
        }

        ///================= 使用哈希值或枚举类型替代字符串比较, 提高性能
        if (dose_type == "average")
        {
            auto voxelVolume_acc = voxelVolume.accessor<double, 1>();
            auto voxelDose_acc = voxelDose.accessor<double, 1>();

            double totalVolume = 0.0;
            double weightedDose = 0.0;
            for (int64_t i = 0; i < voxelVolume.size(0); ++i)
            {
                totalVolume += voxelVolume_acc[i];
                weightedDose += voxelDose_acc[i] * voxelVolume_acc[i];
            }
            if (totalVolume <= 0.0)
            {
                throw std::runtime_error("Prescription ROI volume is zero. Cannot scale to average dose.");
            }

            return torch::Scalar(weightedDose / totalVolume);
        }

        ///================= 提取目标体积百分比
        double targetDoseAtVolume;
        // 使用更高效的查找方式, 例如unordered_map或者switch-case(如果使用枚举)
        if (dose_type == "D_98")
        {
            targetDoseAtVolume = 0.98;
        }
        else if (dose_type == "D_95")
        {
            targetDoseAtVolume = 0.95;
        }
        else if (dose_type == "D_50")
        {
            targetDoseAtVolume = 0.5;
        }
        else if (dose_type == "D_2")
        {
            targetDoseAtVolume = 0.02;
        }
        else if (dose_type == "DoseAtVolume")
        {
            targetDoseAtVolume = doseAtVolume;
        }
        else
        {
            throw std::runtime_error("Prescription type not supported. Only 'average', 'D_98', 'D_95', 'D_50', 'D_2', 'DoseAtVolume' are supported");
        }

        // 调用优化后的get_metric函数
        //auto targetDoseAtVolume_tensor = torch::tensor(targetDoseAtVolume, torch::kFloat64);
        return get_metric(binnedDoseValues, accumulatedRelativeVolumes, targetDoseAtVolume);
    }

}  // namespace dose_torch

///0000000000000000000000000000000000000000000000000000000000000000000000000000000
TORCH_LIBRARY(dose_torch, m)
{
    m.def("accumulate(Tensor a, Tensor b, bool c) -> (Tensor, Tensor)");
    m.def("get_metric(Tensor a, Tensor b, float c) -> Scalar");
    m.def("get_dose_stat_roi_interp(Tensor a, Tensor b, Tensor c, Tensor d, str t, float v) -> Scalar");
}

// 注册PyTorch自定义算子
TORCH_LIBRARY_IMPL(dose_torch, CPU, m)
{
    m.impl("accumulate", &dose_torch::accumulate);
    m.impl("get_metric", &dose_torch::get_metric);
    m.impl("get_dose_stat_roi_interp", &dose_torch::get_dose_stat_roi_interp);
}